"""
    RadianceCostFunction

Module for satellite radiance observation cost function and gradient computation.
Integrates CRTM forward and adjoint operators into variational minimization framework.

# Mathematical Framework

The radiance contribution to the cost function is:
```
J_rad(x) = 0.5 * (H(x) - y)^T * R^{-1} * (H(x) - y)
```

Where:
- H(x): CRTM forward operator mapping atmospheric state to brightness temperatures
- y: Observed brightness temperatures from satellite
- R: Observation error covariance matrix (typically diagonal)

The gradient is computed using the CRTM adjoint operator:
```
∇J_rad(x) = H^T * R^{-1} * (H(x) - y)
```

# Integration with 3D-Var

The radiance cost function is added to the total variational cost:
```
J_total(x) = J_b(x) + J_conv(x) + J_rad(x)
```

Where:
- J_b: Background term
- J_conv: Conventional observation term
- J_rad: Radiance observation term

# Example Usage

```julia
using GSICoreAnalysis
using GSICoreAnalysis.CostFunctions.RadianceCostFunction
using GSICoreAnalysis.DataIO.RadianceBUFRReader
using GSICoreAnalysis.FortranInterface

# Read radiance observations
obs = read_radiance_bufr("gdas1.t12z.1bamua.tm00.bufr_d")

# Initialize CRTM
crtm_init_julia("amsua_n15", "/path/to/coef", obs.channels)

# Create atmospheric state
state = create_atmospheric_state(background)

# Compute radiance cost function
cost = radiance_cost_function(state, obs)

# Compute gradient
grad = radiance_gradient(state, obs)

# Cleanup
crtm_destroy_julia()
```
"""
module RadianceCostFunction

using LinearAlgebra
using Printf
using Statistics

# Import CRTM bindings
include("../FortranInterface/crtm_bindings.jl")

# Import radiance BUFR reader
using ..DataIO: RadianceObservation

export RadianceCostConfig
export AtmosphericProfile, create_atmospheric_profile, profiles_to_crtm_format
export radiance_cost_function, radiance_gradient
export radiance_innovation_statistics, apply_innovation_qc
export setup_radiance_cost_function, cleanup_radiance_cost_function
export simple_forward_operator, simple_adjoint_operator  # For testing

"""
    RadianceCostConfig

Configuration for radiance cost function evaluation.

# Fields
- `use_crtm::Bool`: Use CRTM for radiative transfer (vs simple model)
- `observation_error_scaling::Float64`: Scale factor for observation errors
- `qc_innovation_threshold::Float64`: Innovation threshold for QC (in K)
- `diagonal_r_matrix::Bool`: Use diagonal R matrix (vs full covariance)
- `crtm_initialized::Bool`: Track CRTM initialization status
- `sensor_id::String`: CRTM sensor identifier
- `coef_path::String`: Path to CRTM coefficient files
- `channels::Vector{Int32}`: Active channels
- `apply_channel_mask::Bool`: Honor instrument use flags when selecting channels
"""
mutable struct RadianceCostConfig
    use_crtm::Bool
    observation_error_scaling::Float64
    qc_innovation_threshold::Float64
    diagonal_r_matrix::Bool
    crtm_initialized::Bool
    sensor_id::String
    coef_path::String
    channels::Vector{Int32}
    apply_channel_mask::Bool

    function RadianceCostConfig(;
        use_crtm::Bool = true,
        observation_error_scaling::Float64 = 1.0,
        qc_innovation_threshold::Float64 = 5.0,
        diagonal_r_matrix::Bool = true,
        crtm_initialized::Bool = false,
        sensor_id::String = "",
        coef_path::String = "",
        channels::Vector{Int32} = Int32[],
        apply_channel_mask::Bool = true
    )
        new(use_crtm, observation_error_scaling, qc_innovation_threshold,
            diagonal_r_matrix, crtm_initialized, sensor_id, coef_path, channels,
            apply_channel_mask)
    end
end

function prepare_channel_data(radiance_obs::RadianceObservation,
                              config::RadianceCostConfig)
    channels = isempty(config.channels) ? radiance_obs.channels : config.channels
    obs_bt = radiance_obs.brightness_temp
    obs_err = radiance_obs.obs_error

    if config.apply_channel_mask && isempty(config.channels)
        use_flags = get(radiance_obs.metadata, "channel_use_flags", nothing)
        if use_flags !== nothing && length(channels) == length(use_flags)
            mask = findall(i -> use_flags[i] > 0, eachindex(use_flags))
            if !isempty(mask) && length(mask) < length(channels)
                channels = channels[mask]
                obs_bt = obs_bt[mask, :]
                obs_err = obs_err[mask, :]
            end
        end
    end

    return channels, obs_bt, obs_err
end

"""
    AtmosphericProfile

Atmospheric state profile for radiative transfer calculations.

# Fields
- `temperature::Matrix{Float32}`: Temperature [K] (n_layers × n_profiles)
- `pressure::Matrix{Float32}`: Pressure [hPa] (n_layers × n_profiles)
- `humidity::Matrix{Float32}`: Specific humidity [kg/kg] (n_layers × n_profiles)
- `skin_temperature::Vector{Float32}`: Surface temperature [K] (n_profiles)
- `wind_speed::Vector{Float32}`: Surface wind speed [m/s] (n_profiles)
- `n_layers::Int`: Number of vertical layers
- `n_profiles::Int`: Number of horizontal profiles
"""
struct AtmosphericProfile
    temperature::Matrix{Float32}
    pressure::Matrix{Float32}
    humidity::Matrix{Float32}
    skin_temperature::Vector{Float32}
    wind_speed::Vector{Float32}
    n_layers::Int
    n_profiles::Int

    function AtmosphericProfile(temp, pres, humid, tskin, wspd)
        nlev, nprof = size(temp)
        @assert size(pres) == (nlev, nprof) "Pressure size mismatch"
        @assert size(humid) == (nlev, nprof) "Humidity size mismatch"
        @assert length(tskin) == nprof "Surface temperature size mismatch"
        @assert length(wspd) == nprof "Wind speed size mismatch"
        new(temp, pres, humid, tskin, wspd, nlev, nprof)
    end
end

"""
    create_atmospheric_profile(state_vector::AbstractVector,
                                obs_locations::AbstractMatrix,
                                grid_size::Tuple) -> AtmosphericProfile

Extract atmospheric profiles at observation locations from state vector.

# Arguments
- `state_vector`: Full atmospheric state (flattened)
- `obs_locations`: Observation locations (n_obs × 2: lon, lat)
- `grid_size`: Grid dimensions (nx, ny, nz)

# Returns
- AtmosphericProfile at observation locations
"""
function create_atmospheric_profile(state_vector::AbstractVector{T},
                                    obs_locations::AbstractMatrix,
                                    grid_size::Tuple) where T
    nx, ny, nz = grid_size
    nobs = size(obs_locations, 1)

    # Extract state fields (simplified - assumes specific layout)
    # In real implementation, use proper state vector structure
    n_3d = nx * ny * nz

    # Temperature, humidity, u, v (4 3D fields)
    temp_3d = reshape(state_vector[1:n_3d], nx, ny, nz)
    humid_3d = reshape(state_vector[3*n_3d+1:4*n_3d], nx, ny, nz)

    # Surface pressure (2D field)
    ps_2d = reshape(state_vector[4*n_3d+1:end], nx, ny)

    # Interpolate to observation locations
    # Simplified: nearest neighbor (real implementation would use proper interpolation)
    temp_profiles = zeros(Float32, nz, nobs)
    pres_profiles = zeros(Float32, nz, nobs)
    humid_profiles = zeros(Float32, nz, nobs)
    tskin = zeros(Float32, nobs)
    wspd = zeros(Float32, nobs)

    for i in 1:nobs
        # Get nearest grid point (simplified)
        lon = obs_locations[i, 1]
        lat = obs_locations[i, 2]

        # Map to grid indices (simplified - assumes regular lat/lon grid)
        ix = clamp(round(Int, (lon + 180) / 360 * nx + 1), 1, nx)
        iy = clamp(round(Int, (lat + 90) / 180 * ny + 1), 1, ny)

        # Extract vertical profile
        temp_profiles[:, i] = Float32.(temp_3d[ix, iy, :])
        humid_profiles[:, i] = Float32.(humid_3d[ix, iy, :])

        # Compute pressure levels (simplified - use sigma levels)
        ps = Float32(ps_2d[ix, iy])
        for k in 1:nz
            sigma = Float32(k / (nz + 1))  # Simplified sigma coordinate
            pres_profiles[k, i] = ps * sigma / 100.0  # Convert to hPa
        end

        # Surface properties
        tskin[i] = temp_profiles[1, i]  # Use lowest level temperature
        wspd[i] = 5.0f0  # Default wind speed
    end

    return AtmosphericProfile(temp_profiles, pres_profiles, humid_profiles, tskin, wspd)
end

"""
    profiles_to_crtm_format(profile::AtmosphericProfile) -> Tuple

Convert AtmosphericProfile to CRTM-compatible format.

# Returns
Tuple of (temperature, pressure, humidity, skin_temp, wind_speed) ready for CRTM
"""
function profiles_to_crtm_format(profile::AtmosphericProfile)
    return (profile.temperature, profile.pressure, profile.humidity,
            profile.skin_temperature, profile.wind_speed)
end

"""
    radiance_cost_function(state_vector::AbstractVector,
                           radiance_obs::RadianceObservation,
                           obs_locations::AbstractMatrix,
                           grid_size::Tuple,
                           config::RadianceCostConfig) -> Float64

Compute radiance observation cost function contribution.

# Arguments
- `state_vector`: Atmospheric state vector
- `radiance_obs`: Radiance observations from BUFR
- `obs_locations`: Observation lat/lon locations
- `grid_size`: Model grid dimensions
- `config`: Radiance cost function configuration

# Returns
- Cost function value: J_rad = 0.5 * (H(x) - y)^T * R^{-1} * (H(x) - y)
"""
function radiance_cost_function(state_vector::AbstractVector{T},
                                 radiance_obs::RadianceObservation,
                                 obs_locations::AbstractMatrix,
                                 grid_size::Tuple,
                                 config::RadianceCostConfig) where T

    # Extract atmospheric profiles at observation locations
    profile = create_atmospheric_profile(state_vector, obs_locations, grid_size)

    # Create geometry arrays
    nobs = size(obs_locations, 1)
    zenith = radiance_obs.zenith_angle
    azimuth = radiance_obs.azimuth_angle

    # Determine active channels and associated data
    channels, obs_bt, obs_err = prepare_channel_data(radiance_obs, config)

    # Forward operator: H(x)
    tb_simulated = if config.use_crtm && config.crtm_initialized
        # Use CRTM forward operator
        temp, pres, humid, tskin, wspd = profiles_to_crtm_format(profile)
        crtm_forward_julia(temp, pres, humid, tskin, wspd, zenith, azimuth, channels)
    else
        # Simplified forward operator (for testing without CRTM)
        simple_forward_operator(profile, channels, nobs)
    end

    # Innovation: d = y - H(x)
    innovation = obs_bt .- tb_simulated

    # Apply quality control
    qc_mask = apply_innovation_qc(innovation, obs_err, config)

    # Cost function: J = 0.5 * d^T * R^{-1} * d
    nchans = length(channels)
    cost = 0.0

    if config.diagonal_r_matrix
        # Diagonal R matrix (uncorrelated errors)
        for i in 1:nobs
            for j in 1:nchans
                if qc_mask[j, i]
                    innov = innovation[j, i]
                    obs_err_val = obs_err[j, i] * config.observation_error_scaling

                    # Avoid division by zero or very small errors
                    if obs_err_val > 1e-6
                        cost += 0.5 * (innov^2) / (obs_err_val^2)
                    end
                end
            end
        end
    else
        # Full R matrix (would require covariance structure)
        error("Full R matrix not yet implemented")
    end

    return Float64(cost)
end

"""
    radiance_gradient(state_vector::AbstractVector,
                      radiance_obs::RadianceObservation,
                      obs_locations::AbstractMatrix,
                      grid_size::Tuple,
                      config::RadianceCostConfig) -> Vector

Compute gradient of radiance cost function using CRTM adjoint.

# Arguments
- `state_vector`: Atmospheric state vector
- `radiance_obs`: Radiance observations
- `obs_locations`: Observation locations
- `grid_size`: Model grid dimensions
- `config`: Radiance cost function configuration

# Returns
- Gradient vector: ∇J_rad = H^T * R^{-1} * (H(x) - y)
"""
function radiance_gradient(state_vector::AbstractVector{T},
                           radiance_obs::RadianceObservation,
                           obs_locations::AbstractMatrix,
                           grid_size::Tuple,
                           config::RadianceCostConfig) where T

    # Extract atmospheric profiles at observation locations
    profile = create_atmospheric_profile(state_vector, obs_locations, grid_size)

    # Geometry
    nobs = size(obs_locations, 1)
    zenith = radiance_obs.zenith_angle
    azimuth = radiance_obs.azimuth_angle

    channels, obs_bt, obs_err = prepare_channel_data(radiance_obs, config)

    # Forward operator for innovation
    tb_simulated = if config.use_crtm && config.crtm_initialized
        temp, pres, humid, tskin, wspd = profiles_to_crtm_format(profile)
        crtm_forward_julia(temp, pres, humid, tskin, wspd, zenith, azimuth, channels)
    else
        simple_forward_operator(profile, channels, nobs)
    end

    # Innovation: d = H(x) - y (note sign for gradient)
    innovation = tb_simulated .- obs_bt

    # Apply R^{-1} weighting
    nchans = length(channels)
    weighted_innovation = zeros(Float32, nchans, nobs)
    qc_mask = apply_innovation_qc(innovation, obs_err, config)

    for i in 1:nobs
        for j in 1:nchans
            if qc_mask[j, i]
                obs_err_val = obs_err[j, i] * config.observation_error_scaling
                if obs_err_val > 1e-6
                    weighted_innovation[j, i] = innovation[j, i] / (obs_err_val^2)
                end
            end
        end
    end

    # Adjoint operator: H^T * (R^{-1} * innovation)
    gradient_profiles = if config.use_crtm && config.crtm_initialized
        # Use CRTM adjoint operator
        temp, pres, humid, tskin, wspd = profiles_to_crtm_format(profile)
        temp_ad, humid_ad = crtm_adjoint_julia(weighted_innovation, temp, pres, humid,
                                               tskin, wspd, zenith, azimuth, channels)
        (temp_ad, humid_ad)
    else
        # Simplified adjoint operator
        simple_adjoint_operator(weighted_innovation, profile, channels, nobs)
    end

    # Map gradient from observation space back to state vector space
    gradient = map_gradient_to_state_vector(gradient_profiles, obs_locations, grid_size)

    return gradient
end

"""
    apply_innovation_qc(innovation::Matrix, obs_error::Matrix, config::RadianceCostConfig) -> BitMatrix

Apply quality control based on innovation statistics.

# Returns
Boolean mask indicating which observations pass QC (true = good)
"""
function apply_innovation_qc(innovation::Matrix{T}, obs_error::Matrix{T},
                             config::RadianceCostConfig) where T
    nchans, nobs = size(innovation)
    qc_mask = trues(nchans, nobs)

    for i in 1:nobs
        for j in 1:nchans
            # Check if innovation is too large (O-B check)
            normalized_innov = abs(innovation[j, i]) / obs_error[j, i]

            if normalized_innov > config.qc_innovation_threshold
                qc_mask[j, i] = false
            end

            # Check for NaN or missing values
            if isnan(innovation[j, i]) || isnan(obs_error[j, i])
                qc_mask[j, i] = false
            end
        end
    end

    return qc_mask
end

"""
    map_gradient_to_state_vector(gradient_profiles::Tuple,
                                  obs_locations::AbstractMatrix,
                                  grid_size::Tuple) -> Vector

Map gradient from observation locations back to full state vector grid.

# Arguments
- `gradient_profiles`: (temp_gradient, humid_gradient) at obs locations
- `obs_locations`: Observation locations
- `grid_size`: Model grid dimensions

# Returns
- Full state vector gradient
"""
function map_gradient_to_state_vector(gradient_profiles::Tuple,
                                      obs_locations::AbstractMatrix,
                                      grid_size::Tuple)
    temp_grad, humid_grad = gradient_profiles
    nx, ny, nz = grid_size
    nobs = size(obs_locations, 1)

    # Initialize gradient arrays
    n_3d = nx * ny * nz
    n_2d = nx * ny
    n_state = 4 * n_3d + n_2d  # 4 3D fields + 1 2D field

    gradient = zeros(n_state)

    # Temperature gradient contribution
    temp_grad_3d = zeros(nx, ny, nz)
    humid_grad_3d = zeros(nx, ny, nz)

    # Distribute gradients from observation locations to grid
    for i in 1:nobs
        lon = obs_locations[i, 1]
        lat = obs_locations[i, 2]

        # Map to grid indices (simplified)
        ix = clamp(round(Int, (lon + 180) / 360 * nx + 1), 1, nx)
        iy = clamp(round(Int, (lat + 90) / 180 * ny + 1), 1, ny)

        # Accumulate gradient at grid point
        for k in 1:nz
            temp_grad_3d[ix, iy, k] += temp_grad[k, i]
            humid_grad_3d[ix, iy, k] += humid_grad[k, i]
        end
    end

    # Pack into state vector gradient
    gradient[1:n_3d] = vec(temp_grad_3d)
    gradient[3*n_3d+1:4*n_3d] = vec(humid_grad_3d)
    # u, v, ps gradients remain zero (radiance not sensitive to winds/surface pressure)

    return gradient
end

"""
    simple_forward_operator(profile::AtmosphericProfile,
                            channels::Vector{Int32},
                            nobs::Int) -> Matrix{Float32}

Simplified forward operator for testing (when CRTM not available).
Uses simple weighting function approximation.
"""
function simple_forward_operator(profile::AtmosphericProfile,
                                 channels::Vector{Int32},
                                 nobs::Int)
    nchans = length(channels)
    nlev = profile.n_layers
    tb_simulated = zeros(Float32, nchans, nobs)

    # Simple weighting function model
    for i in 1:nobs
        for (j, ch) in enumerate(channels)
            # Different channels peak at different levels
            peak_level = clamp(Int(round(ch * nlev / 15)), 1, nlev)

            # Weighted average of temperature
            total_weight = 0.0f0
            weighted_temp = 0.0f0

            for k in 1:nlev
                # Gaussian-like weighting function
                weight = Float32(exp(-(k - peak_level)^2 / (nlev/4)^2))
                weighted_temp += weight * profile.temperature[k, i]
                total_weight += weight
            end

            tb_simulated[j, i] = weighted_temp / total_weight
        end
    end

    return tb_simulated
end

"""
    simple_adjoint_operator(weighted_innovation::Matrix{Float32},
                            profile::AtmosphericProfile,
                            channels::Vector{Int32},
                            nobs::Int) -> Tuple

Simplified adjoint operator (transpose of simple forward operator).
"""
function simple_adjoint_operator(weighted_innovation::Matrix{Float32},
                                 profile::AtmosphericProfile,
                                 channels::Vector{Int32},
                                 nobs::Int)
    nchans = length(channels)
    nlev = profile.n_layers

    temp_ad = zeros(Float32, nlev, nobs)
    humid_ad = zeros(Float32, nlev, nobs)

    # Adjoint of weighted average
    for i in 1:nobs
        for (j, ch) in enumerate(channels)
            peak_level = clamp(Int(round(ch * nlev / 15)), 1, nlev)

            total_weight = 0.0f0
            for k in 1:nlev
                weight = Float32(exp(-(k - peak_level)^2 / (nlev/4)^2))
                total_weight += weight
            end

            # Distribute innovation to levels
            for k in 1:nlev
                weight = Float32(exp(-(k - peak_level)^2 / (nlev/4)^2))
                temp_ad[k, i] += (weight / total_weight) * weighted_innovation[j, i]
            end
        end
    end

    return (temp_ad, humid_ad)
end

"""
    radiance_innovation_statistics(innovation::Matrix, obs_error::Matrix, qc_mask::BitMatrix) -> Dict

Compute statistics of radiance innovations for monitoring.

# Returns
Dictionary with statistics per channel
"""
function radiance_innovation_statistics(innovation::Matrix{T},
                                        obs_error::Matrix{T},
                                        qc_mask::BitMatrix) where T
    nchans, nobs = size(innovation)
    stats = Dict{String,Any}()

    for j in 1:nchans
        good_obs = [innovation[j, i] for i in 1:nobs if qc_mask[j, i]]

        if length(good_obs) > 0
            stats["channel_$j"] = Dict(
                "mean" => mean(good_obs),
                "std" => std(good_obs),
                "rms" => sqrt(mean(good_obs.^2)),
                "min" => minimum(good_obs),
                "max" => maximum(good_obs),
                "n_good" => length(good_obs),
                "n_total" => nobs
            )
        end
    end

    return stats
end

"""
    setup_radiance_cost_function(sensor_id::String, coef_path::String,
                                  channels::Vector{Int32}) -> RadianceCostConfig

Initialize CRTM and setup radiance cost function configuration.

# Arguments
- `sensor_id`: CRTM sensor identifier (e.g., "amsua_n15")
- `coef_path`: Path to CRTM coefficient files
- `channels`: Channels to use

# Returns
- Configured RadianceCostConfig
"""
function setup_radiance_cost_function(sensor_id::String, coef_path::String,
                                      channels::Vector{Int32})
    config = RadianceCostConfig(
        use_crtm = true,
        sensor_id = sensor_id,
        coef_path = coef_path,
        channels = channels
    )

    # Initialize CRTM
    status = crtm_init_julia(sensor_id, coef_path, channels)

    if status == 0
        config.crtm_initialized = true
        @info "CRTM initialized successfully for $sensor_id"
    else
        @warn "CRTM initialization failed with status $status, using simplified operator"
        config.use_crtm = false
        config.crtm_initialized = false
    end

    return config
end

"""
    cleanup_radiance_cost_function(config::RadianceCostConfig)

Cleanup CRTM resources.
"""
function cleanup_radiance_cost_function(config::RadianceCostConfig)
    if config.crtm_initialized
        crtm_destroy_julia()
        config.crtm_initialized = false
        @info "CRTM cleaned up"
    end
end

end # module RadianceCostFunction
